import numpy as np
from collections import defaultdict


def calculate_distance_vector(point_a, point_b):
    """
    计算两个点之间的欧几里得距离
    """
    diff_squared = [(a - b) ** 2 for a, b in zip(point_a, point_b)]
    return sum(diff_squared) ** 0.5


class KNearestNeighborsClassifier:
    """
    K最近邻分类器实现
    """

    def __init__(self, neighbor_count=3, category_count=2):
        """
        初始化分类器参数
        """
        if neighbor_count <= 0:
            raise ValueError("邻居数量必须为正整数")

        self.neighbor_count = neighbor_count
        self.category_count = category_count
        self.training_features = None
        self.training_labels = None

    def train_model(self, feature_data, label_data):
        """
        训练模型：存储训练数据
        """
        if len(feature_data) != len(label_data):
            raise ValueError("特征和标签数量不匹配")

        self.training_features = feature_data
        self.training_labels = label_data

    def _find_nearest_neighbors(self, test_point):
        """
        寻找测试点的最近邻居索引
        """
        distance_calculations = []

        for idx, train_point in enumerate(self.training_features):
            dist = calculate_distance_vector(test_point, train_point)
            distance_calculations.append((dist, idx))

        # 按距离排序并选择前k个
        distance_calculations.sort(key=lambda x: x)
        neighbor_indices = [item[1] for item in distance_calculations[:self.neighbor_count]]

        return neighbor_indices

    def _determine_class_label(self, neighbor_indices):
        """
        根据邻居的标签确定测试点的类别
        """
        label_frequency = defaultdict(int)

        for idx in neighbor_indices:
            current_label = self.training_labels[idx]
            label_frequency[current_label] += 1

        # 选择出现频率最高的标签
        max_frequency = -1
        selected_label = None

        for label, count in label_frequency.items():
            if count > max_frequency or (count == max_frequency and label < selected_label):
                max_frequency = count
                selected_label = label

        return selected_label

    def classify_points(self, test_data):
        """
        对测试数据进行分类预测
        """
        if self.training_features is None:
            raise ValueError("模型尚未训练，请先调用train_model方法")

        prediction_results = []

        for test_point in test_data:
            # 步骤1: 找到最近邻居
            neighbor_indices = self._find_nearest_neighbors(test_point)

            # 步骤2: 投票决定类别
            predicted_label = self._determine_class_label(neighbor_indices)

            prediction_results.append(predicted_label)

        return prediction_results


# ==================== 演示示例 ====================

def demonstrate_classification():
    """
    演示KNN分类器的使用
    """
    # 准备训练数据
    training_features = [
        [1.0, 2.0], [1.5, 1.8], [5.0, 8.0],
        [6.0, 9.0], [1.0, 0.6], [9.0, 11.0]
    ]

    training_labels = [0, 0, 1, 1, 0, 1]

    # 准备测试数据
    test_features = [
        [2.0, 3.0],
        [8.0, 9.0],
        [0.5, 0.5]
    ]

    # 创建并训练分类器
    classifier = KNearestNeighborsClassifier(neighbor_count=3, category_count=2)
    classifier.train_model(training_features, training_labels)

    # 进行预测
    predictions = classifier.classify_points(test_features)

    # 输出结果
    print("测试数据分类结果:")
    for i, (test_point, pred) in enumerate(zip(test_features, predictions)):
        print(f"测试点 {test_point} → 预测类别: {pred}")

    return predictions


if __name__ == "__main__":
    result_predictions = demonstrate_classification()
    print(f"\n最终预测结果列表: {result_predictions}")
